use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tonic::codegen::tokio_stream::Stream;
use tonic::{Request, Response, Status};
use tower::Layer;
use tracing::{info, Level};

pub mod auth_service {
    tonic::include_proto!("auth_service");
}

pub mod user_service {
    tonic::include_proto!("user_service");
}

use auth_service::{
    auth_service_server::{AuthService, AuthServiceServer},
    LoginRequest, LoginResponse, VerifyTokenRequest, VerifyTokenResponse
};

use user_service::{
    user_service_server::{UserService, UserServiceServer}, CreateUserRequest, CreateUserResponse, GetUserRequest,
    GetUserResponse, ListUsersRequest, User,
};

#[derive(Debug)]
pub struct ConnectionStats {
    pub total_connections: usize,
    pub broken_pipes: usize,
    pub successful_reconnects: usize,
    pub active_connections: usize,
}

impl ConnectionStats {
    pub fn new() -> Self {
        Self {
            total_connections: 0,
            broken_pipes: 0,
            successful_reconnects: 0,
            active_connections: 0,
        }
    }
}

#[derive(Debug)]
pub struct MyUserService {
    stats: Arc<RwLock<ConnectionStats>>,
    server_id: String,
}

impl MyUserService {
    pub fn new() -> Self {
        Self {
            stats: Arc::new(RwLock::new(ConnectionStats::new())),
            server_id: format!("server-{}", rand::random::<u16>()),
        }
    }
}

#[tonic::async_trait]
impl UserService for MyUserService {
    async fn get_user(
        &self,
        request: Request<GetUserRequest>,
    ) -> Result<Response<GetUserResponse>, Status> {
        let request = request.into_inner();
        let user_id = request.user_id;

        let user = User {
            id: user_id,
            name: "WUZHIJUN".to_string(),
            email: "wu_zhijun@msn.com".to_string(),
            created_at: chrono::Utc::now().timestamp_millis(),
        };

        {
            let mut stats = self.stats.write().await;
            stats.total_connections += 1;
            stats.active_connections += 1;
        }

        let response = GetUserResponse { user: Some(user) };

        Ok(Response::new(response))
    }

    async fn create_user(
        &self,
        _request: Request<CreateUserRequest>,
    ) -> Result<Response<CreateUserResponse>, Status> {
        let user = User {
            id: "wu_zhijun".to_string(),
            name: "WUZHIJUN".to_string(),
            email: "wu_zhijun@msn.com".to_string(),
            created_at: chrono::Utc::now().timestamp_millis(),
        };

        let response = CreateUserResponse { user: Some(user) };

        Ok(Response::new(response))
    }

    type ListUsersStream = Pin<Box<dyn Stream<Item = Result<User, Status>> + Send + 'static>>;

    async fn list_users(
        &self,
        _request: Request<ListUsersRequest>,
    ) -> Result<Response<Self::ListUsersStream>, Status> {
        let stream = async_stream::stream! {
            for i in 0..5 {
                let user = User {
                    id: format!("wu_zhijun-{}", i),
                    name: format!("WUZHIJUN-{}", i),
                    email: format!("wu_zhijun-{}@msn.com", i),
                    created_at: chrono::Utc::now().timestamp_millis(),
                };

                yield Ok(user);

                tokio::time::sleep(Duration::from_millis(200)).await;
            }
        };
        Ok(Response::new(Box::pin(stream)))
    }
}

#[derive(Debug)]
pub struct MyAuthService {
    stats: Arc<RwLock<ConnectionStats>>,
    server_id: String,
}

impl MyAuthService {
    pub fn new() -> Self {
        Self {
            stats: Arc::new(RwLock::new(ConnectionStats::new())),
            server_id: format!("server-{}", rand::random::<u16>()),
        }
    }
}

#[tonic::async_trait]
impl AuthService for MyAuthService {
    async fn login(&self, _request: Request<LoginRequest>) -> Result<Response<LoginResponse>, Status> {
        let response = LoginResponse {
            token: "a1334324324".to_string(),
            expires_at: 1234,
        };
        Ok(Response::new(response))
    }

    async fn verify_token(&self, _request: Request<VerifyTokenRequest>) -> Result<Response<VerifyTokenResponse>, Status> {
        let response = VerifyTokenResponse {
            valid: true,
            user_id: "wu_zhijun".to_string(),
        };
        Ok(Response::new(response))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt().with_max_level(Level::DEBUG).init();

    let addr = "[::1]:50051".parse()?;

    let user_service = MyUserService::new();
    let auth_service = MyAuthService::new();

    let stats = Arc::clone(&user_service.stats);

    info!("gRPC服务器启动，监听地址: {}", addr);
    info!("服务器ID: {}", user_service.server_id);

    // 启动统计信息报告
    let stats_clone = Arc::clone(&stats);
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(10)).await;

            let stats = stats_clone.read().await;
            info!("服务器统计 - 总连接: {}, Broken pipes: {}, 成功重连: {}, 活跃连接: {}", stats.total_connections,
                stats.broken_pipes,
                stats.successful_reconnects,
                stats.active_connections
            );
        }
    });

    tonic::transport::Server::builder()
        .add_service(UserServiceServer::new(user_service))
        .add_service(AuthServiceServer::new(auth_service))
        .serve(addr)
        .await?;

    Ok(())
}
